from pyson.types import is_container
from pyson.iface.iterators import iteritems

__all__ = ['walkitems', 'walkvalues', 'walkpaths', 'walkequal', 'walknotequal',
           'walkdiff', 'zipwalk']
_inf = float('inf')

def walkitems(obj, anyiter=False, ret_nodes=False, max_nest=_inf, trunc=None):
    '''
    Iterate over all pairs of (path, values) in obj. 
    
    If 'obj' is a mapping, iteritems(obj) is equivalent obj.iteritems(). If it 
    is a sequence, iteritems(obj) behaves similarly to enumerate(obj).
    
    Parameters
    ----------
    
    obj : any container object (mapping or sequence)
        Object to be iterated
    anyiter : bool
        If 'obj' is iterable but does not support indexing (e.g., an iterator 
        or a set object), the default behavior is to raise a TypeError. If 
        'anyiter' is True, this function works similarly to the enumerate(obj) 
        for these objects.
    ret_nodes : bool
        If True, it will also yield (path, value) pairs for container objects
    max_nest : int
        The maximum nesting in which to explore the object
    trunc : bool or None
        Control the behavior when maximum nesting is reached. See examples 
        bellow.
    
    Examples
    --------
    
    Given the structure
    
    >>> json = {'foo': 'bar',
    ...         'ham': {'spam': True, 'eggs': [False, None]}}
    
    It will iterate over all paths and values in the structure
    
    >>> from pyson import *
    >>> printiter(walkitems(json))
    (('foo',), 'bar')
    (('ham', 'spam'), True)
    (('ham', 'eggs', 0), False)
    (('ham', 'eggs', 1), None)
    
    Nesting can be controlled. This is useful for dealing with very nested
    (possibly infinitely nested) structures:
    
    >>> json = {
    ...     'nesting': 0,
    ...     'nested': {
    ...         'nesting': 1, 
    ...         'sub-nested': {
    ...             'nesting': 2, 
    ...             'sub-sub-nested': { 'nesting': 3, 'value': 'Hello!'}
    ...         }
    ...     }
    ... }
    >>> printiter(walkitems(json, max_nest=2))
    (('nesting',), 0)
    (('nested', 'nesting'), 1)
    (('nested', 'sub-nested', 'nesting'), 2)
    (('nested', 'sub-nested', 'sub-sub-nested'), {'nesting': 3, 'value': 'Hello!'})
    
    When the maximum nesting is reached, container objects are returned as if 
    they are values (thats why the full dictionary was returned in the last 
    line)
    
    One can simply discard the values above a certain nesting by setting the
    optional argument ``trunc`` to a True value:
    
    >>> printiter(walkitems(json, max_nest=2, trunc=True))
    (('nesting',), 0)
    (('nested', 'nesting'), 1)
    (('nested', 'sub-nested', 'nesting'), 2)

    If ``trunc`` is False (except for None), it will raise a ValueError
    
    >>> printiter(walkitems(json, max_nest=2, trunc=False))
    Traceback (most recent call last):
    ...
    ValueError: maximum nesting reached
    
    Finally, one can return the container sub-nodes by setting ``ret_nodes`` to
    True
    
    >>> json = {'foo': 'bar',
    ...         'ham': {'spam': True, 'eggs': False}}
    >>> printiter(walkitems(json, ret_nodes=True))
    (('foo',), 'bar')
    (('ham',), {'eggs': False, 'spam': True})
    (('ham', 'eggs'), False)
    (('ham', 'spam'), True)
    
    A zero nesting iteration is similar to iteritems(), but return (path, value) 
    items instead of (key, value) pairs.
    
    >>> list(walkitems({0: 'zero', 1: [1, 2], 42: 'answer'}, max_nest=0))
    [((0,), 'zero'), ((1,), [1, 2]), ((42,), 'answer')]
    
    See also
    --------
        
    walkpaths()  --- iterate over all paths of 'obj'
    walkvalues() --- iterate over all values of 'obj'
    '''

    if not is_container(obj):
        raise ValueError('expected a container object')

    stack = [(0, (), obj)]
    while stack:
        nesting, path, obj = stack.pop()
        if is_container(obj):
            for key, value in iteritems(obj, anyiter):
                subkey = path + (key,)
                if is_container(value):
                    if ret_nodes:
                        yield (subkey, value)

                    nesting += 1
                    if nesting > max_nest:
                        if trunc is None:
                            if not ret_nodes:
                                yield (subkey, value)
                        elif trunc:
                            pass
                        else:
                            raise ValueError('maximum nesting reached')
                    else:
                        stack.append((nesting, subkey, value))
                else:
                    yield (subkey, value)
        else:
            yield (path, obj)

# Functions based on iteritems() -----------------------------------------------
def walkvalues(obj, anyiter=False, ret_nodes=False, max_nest=_inf, trunc=None):
    '''
    Iterate over all values of 'obj'.
    
    See also
    --------
    
    walkitems()  --- iterate over (key, obj[key]) pairs
    walkpaths()  --- iterate over all paths of 'obj'
    '''
    for (_, v) in walkitems(obj, anyiter, ret_nodes, max_nest, trunc):
        yield v

def walkpaths(obj, anyiter=False, ret_nodes=False, max_nest=_inf, trunc=None):
    '''
    Iterate over all keys/indices of 'obj'
    
    See also
    --------
    
    walkitems()  --- iterate over (key, obj[key]) pairs
    walkvalues() --- iterate over all values of 'obj'
    '''
    for (p, _) in walkitems(obj, anyiter, ret_nodes, max_nest, trunc):
        yield p

def walkequal(x, y):
    '''Iterate over all values (path, value) pairs that are identical between x 
    and y.'''

    flat = dict(walkitems(x))
    flat_other = dict(walkitems(y))
    for k, v in flat.items():
        try:
            if v == flat_other[k]:
                yield (k, v)
        except KeyError:
            continue

def walknotequal(x, y, force_common=False):
    '''Iterate over all values (path, value) pairs in ``x`` that are different 
    between x and y.'''

    flat = dict(walkitems(x))
    flat_other = dict(walkitems(y))
    for k, v in flat.items():
        try:
            if v != flat_other.pop(k):
                yield (k, v)
        except KeyError:
            if not force_common:
                yield (k, v)

def walkdiff(x, y, force_common=False, undefined=None):
    '''Iterate over all values (path, value) pairs that are different in ``x`` and
    ``y`` and return (path, (value_x, value_y)) tuples.'''

    flat = dict(walkitems(x))
    flat_other = dict(walkitems(y))
    for k, v in flat.items():
        try:
            v_other = flat_other.pop(k)
            if v != v_other:
                yield (k, (v, v_other))
        except KeyError:
            if not force_common:
                yield (k, (v, undefined))
    if not force_common:
        for k, v in flat_other.items():
            yield (k, (undefined, v))


def zipwalk(*args, **kwds):
    '''Iterate over all tuples of values in the arguments.
    
    If the optional keyword argument ``paths`` is True, it will return 
    a tuple with (path, v1, v2, ...).'''

    # Get keywords
    ret_keys = kwds.pop('paths', False)
    if kwds:
        raise TypeError('unsupported keyword argument: %s' % kwds.pop())

    # Special case len 0 and 1 args
    if len(args) == 0:
        return
    elif len(args) == 1:
        dispatcher = {True: walkitems, False:walkvalues}
        worker = dispatcher[ret_keys]
        for x in worker(args[0]):
            yield x

    # General case
    flats = [dict(walkitems(x)) for x in args[1:]]
    base_flat = dict(walkitems(args[0]))

    for k, v in base_flat.items():
        result = [k, v] if ret_keys else [v]
        for flat in flats:
            try:
                result.append(flat[k])
            except KeyError:
                break
        else:
            yield tuple(result)

if __name__ == '__main__':
    import doctest
    doctest.testmod(optionflags=doctest.REPORT_ONLY_FIRST_FAILURE, verbose=0)

